import math

try :    
    cmds.delete( dellist )
    dellist = []
except :
    dellist = []

def distance(a,b):
    posA = cmds.xform( a , q=1 , rp=1 )
    posB = cmds.xform( b , q=1 , rp=1 )
    tot = 0
    for i in range(3):
        tot +=(posA[i]-posB[i])**2
    return math.sqrt(tot)

def getHich( tg ):    
    main = [ m for m in cmds.listRelatives( tg,c = 1) if 'main' in m]
    leg_grp = [ l for l in cmds.listRelatives( tg,c = 1) if 'leg_grp' in l]
    legs = cmds.listRelatives( leg_grp[0],c = 1)
    return {'main':main , 'legs':legs}
    
mode = 1
radL = 0
radR = 0

def antwalk( ant ):
    global radL
    global radR
    ct = cmds.currentTime( q=1 )
    
    ant_main = ant['main'][0]
    ant_lfft_ctrl = [a for a in ant['legs'] if 'lfft' in a ][0]
    ant_lfmd_ctrl = [a for a in ant['legs'] if 'lfmd' in a ][0]
    ant_lfbk_ctrl = [a for a in ant['legs'] if 'lfbk' in a ][0]
    ant_rtft_ctrl = [a for a in ant['legs'] if 'rtft' in a ][0]
    ant_rtmd_ctrl = [a for a in ant['legs'] if 'rtmd' in a ][0]
    ant_rtbk_ctrl = [a for a in ant['legs'] if 'rtbk' in a ][0]
    
    
    # initial
    if ct == 1 :
        for x in ant['legs']:
            map( lambda s :cmds.setAttr( '%s.t' % x + s , 0 ) , 'xy' )
               
        
    mainX = cmds.getAttr( ant_main + '.tx' )
    mainX1 = cmds.getAttr( ant_main + '.tx' , t = ct-1 )
    mainT = cmds.getAttr( ant_main + '.t' , t = ct-1 ) 
    dist2Main = mainX - mainX1

    upSpeed = 5.0
    liftH = .15
    
    ## Front

    if not dist2Main and not max( 0 , math.sin( radR ) * liftH ) :
        pass
    else:
        if cmds.getAttr( ant_lfft_ctrl + '.ty') == 0 :
            cmds.setAttr( ant_rtft_ctrl + '.ty' , max( 0 , math.sin( radR ) * liftH ) )       
            radR = ( radR + math.pi / upSpeed )        

    if cmds.getAttr(ant_rtft_ctrl+'.ty') != 0 :        
        cmds.setAttr( ant_rtft_ctrl+'.tx' , mainX ) 
    
    if not dist2Main and not max( 0 , math.sin( radL ) * liftH ) :
        pass
    else:
        if cmds.getAttr(ant_rtft_ctrl+'.ty') == 0 :
            cmds.setAttr( ant_lfft_ctrl+'.ty' , max( 0 , math.sin( radL ) * liftH ) )        
            radL = ( radL + math.pi / upSpeed )
            
    if cmds.getAttr(ant_lfft_ctrl+'.ty') != 0 :
        cmds.setAttr( ant_lfft_ctrl+'.tx' , mainX )
        
        
    ## Middle
    if not dist2Main and not max( 0 , math.sin( radR ) * liftH ) :
        pass
    else:
        if cmds.getAttr(ant_lfmd_ctrl+'.ty') == 0 :
            cmds.setAttr( ant_rtmd_ctrl+'.ty' , max( 0 , math.sin( radR ) * liftH ) )        
            radR = ( radR + math.pi / upSpeed )
                
    if cmds.getAttr(ant_rtmd_ctrl+'.ty') != 0 :
        cmds.setAttr( ant_rtmd_ctrl+'.tx' , mainX )        
    
    if not dist2Main and not max( 0 , math.sin( radL ) * liftH ) :
        pass
    else:
        if cmds.getAttr(ant_rtmd_ctrl+'.ty') == 0 :        
            cmds.setAttr( ant_lfmd_ctrl+'.ty' , max( 0 , math.sin( radL ) * liftH ) )        
            radL = ( radL + math.pi / upSpeed )
                
        
    if cmds.getAttr(ant_lfmd_ctrl+'.ty') != 0 :
        cmds.setAttr( ant_lfmd_ctrl+'.tx' , mainX )
        
        
        
    ## Back
    if not dist2Main and not max( 0 , math.sin( radR ) * liftH ) :
        pass
    else: 
        if cmds.getAttr(ant_lfbk_ctrl+'.ty') == 0 :
            cmds.setAttr( ant_rtbk_ctrl+'.ty' , max( 0 , math.sin( radR ) *  liftH ) )
            radR = ( radR + math.pi / upSpeed ) 
    
        
    if cmds.getAttr(ant_rtbk_ctrl+'.ty') != 0 :
        cmds.setAttr( ant_rtbk_ctrl+'.tx' , mainX )        
    
    
    if  dist2Main == 0  and  max( 0 , math.sin( radL ) * liftH ) ==0 :
        pass
    else:
        if cmds.getAttr(ant_rtbk_ctrl+'.ty') == 0 :        
            cmds.setAttr( ant_lfbk_ctrl+'.ty' , max( 0 , math.sin( radL ) * liftH ) )
            radL = ( radL + math.pi / upSpeed )        
        
    if cmds.getAttr(ant_lfbk_ctrl+'.ty') != 0 :
        cmds.setAttr( ant_lfbk_ctrl+'.tx' , mainX )
        
    if  dist2Main == 0  :
        for x in ant['legs']:
            cmds.setAttr( '%s.ty' % x  , 0 )




    
for x in range( cmds.playbackOptions(q=1,aet=1)):
    cmds.currentTime(x+1)
    for iAnt in cmds.ls('ant_all*' , tr=1):
        ant = getHich(iAnt)
        antwalk(ant)
        for lg in ant['legs']:
            map( lambda s :cmds.setKeyframe( lg ,at='translate'+s ) , 'XY' )
        


#exp = cmds.expression( s='python("antwalk()")' , ae = 1 )
#dellist.append( exp )




for iAnt in cmds.ls('ant_all*' , tr=1):
    ant = getHich(iAnt)
    for x in range( cmds.playbackOptions(q=1,aet=1)):
        cmds.currentTime(x+1) 
        antwalk(ant)
        for lg in ant['legs']:
            map( lambda s :cmds.setKeyframe( lg ,at='translate'+s ) , 'XY' )


